package com.qzdatasoft.mcp.server.common;

import lombok.RequiredArgsConstructor;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;

import java.util.Map;
import java.util.Optional;

@Aspect
@Component
@RequiredArgsConstructor
public class JdbcTemplateInterceptor {

    private JdbcPermission jdbcPermission;

    @Before("execution(* org.springframework.jdbc.core.JdbcTemplate.*(java.lang.String, ..))")
    public void beforeExecute(JoinPoint joinPoint) {
        Object[] args = joinPoint.getArgs();
        for (Object arg : args) {
            if (arg instanceof String sql) {
                if(!ObjectUtils.isEmpty(jdbcPermission)) jdbcPermission.check(sql);
                break;
            }
        }
    }

    public interface JdbcPermission {
        Map<String, String[]> getRoles();

        String getCurrentDbType();

        JdbcTemplate getJdbcTemplate();

        default String[] getPermissions(){
            return getRoles().get(getCurrentRole());
        }

        default void check(String sql)
                throws ExecuteSqlNoPermissionException {
            String[] permissions = this.getPermissions();
            for (String permission : permissions) {
                if (sql.trim().toUpperCase().startsWith(permission)) {
                    return;
                }
            }
            throw new ExecuteSqlNoPermissionException("""
Insufficient permissions: The current role \\"%s\\" does not have permission to perform this SQL operation"""
                    .formatted(getCurrentRole()));
        }

        default String getCurrentRole(){
            return Optional.ofNullable((String)RequestThreadLocal.get("MCP_DB_ROLE")).orElse("readonly");
        }

    }
}